Allocating workspaces fails due to invalid view call

WorkspaceManager.get reinterprets the byte buffer with current_workspace[:num_bytes].view(spec.dtype) but Tensor.view only accepts a shape, not a dtype. Passing a torch.dtype raises TypeError: 'torch.dtype' object cannot be interpreted as an integer, so every call to reserve/get will crash before returning a workspace. The manager needs to reshape using a size tuple and cast with view(dtype) via reinterpret_cast semantics (e.g. view(-1).view(spec.dtype) or view(dtype).reshape).

Useful? React with 👍 / 👎.

Prefill-only batches reference attn_out before initialization

In the fp8 path of FlashMLASparseImpl.forward, attn_out is only assigned inside the if num_decode_tokens > 0 branch. The subsequent if num_prefill_tokens > 0 branch unconditionally reads decode_attn_out = attn_out, which raises UnboundLocalError whenever a batch contains only prefill tokens. Prefill batches are common during initial context ingestion, so this path will always fail until attn_out is initialized for the prefill case.

Useful? React with 👍 / 👎.

Decode outputs stored into prefill slots

When both decode and prefill tokens exist, the fp8 path copies decode attention results with attn_out[:num_prefill_tokens] = decode_attn_out[:num_prefill_tokens]. Decode tokens occupy the first num_decode_tokens entries, so this writes them into the wrong slice and fails whenever num_prefill_tokens > num_decode_tokens because the right-hand side is shorter than the target. The assignment should use num_decode_tokens to preserve decode outputs and avoid size mismatches.

Useful? React with 👍 / 👎.

do we have to do these allocation during model execution? Is it possible to setup the memory buffer before real execution to reduce the runtime overhead?
For example, I'm thinking of

1. during model init, call current_workspace_manager().get_simultaneous() to tell the workspace manager the max possible size the model may used
2. lock the memory space
3. during profile run, allocate the memory and save to self.xxx like self.workspace13 = workspace13; self.workspace2 = workspace2
4. during model execution, just use the self.workspace13

ya we should do something like this the only complication currently is when self.fused_experts.supports_chunking() == False (i.e. PPLX or DeepEP LL) then we need profile run to know the shape of the workspaces because in that case the profile run actually mirrors the worst case scenario (hence gating this logic).

This code is meant to a straight refactor of #27426 (once that lands) to use the new workspace manager so im partial to leaving this optimization to a future PR if you are cool with it. I am trying to learn more about the MoE chunking code in-order to propose a broader UX refactor there since I find it confusing currently and I think this optimization could be part of that.

Yeah works for me. Can you do some benchmark to ensure no perf regression?

should these lines in sparse_attn_indexer_fake be removed? I wrote them to mimic the activation memory during profile run.

_flattened_kv = torch.empty(
        [total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8
    )
    _k_fp8 = _flattened_kv[..., :head_dim].view(torch.float8_e4m3fn).contiguous()
    _k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous()

theyre back 👍 (not sure when haha; I guess I must have done it and just not pushed it)

I want to say that if you choose to reserve the buffers of shape (total_seq_lens, head_dim) and (total_seq_lens, 4) in the workspace, you don't need to run the following in sparse_attn_indexer_fake

_flattened_kv = torch.empty(
        [total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8
    )

to also reserve memory for these two tensors in activation memory.